// Copyright 2022 Chen Jun
// Licensed under the MIT License.

#ifndef ARMOR_DETECTOR__PNP_SOLVER_HPP_
#define ARMOR_DETECTOR__PNP_SOLVER_HPP_

#include <geometry_msgs/msg/point.hpp>

// STD
#include <array>
#include <vector>

#include <opencv2/calib3d.hpp>

#include "armor_detector/armor.hpp"

namespace rm_auto_aim
{
class PnPSolver
{
public:
  PnPSolver(
    const std::array<double, 9> & camera_matrix,
    const std::vector<double> & distortion_coefficients);

  /**
   * @brief 解算装甲板pnp
   * 
   * @param armor 装甲板信息
   * @param rvec 旋转向量
   * @param tvec 平移向量
   * @return true 
   * @return false 
   */
  bool solvePnP(const Armor & armor, cv::Mat & rvec, cv::Mat & tvec);

  // Calculate the distance between armor center and image center
  float calculateDistanceToCenter(const cv::Point2f & image_point);

private:
  cv::Mat camera_matrix_;
  cv::Mat dist_coeffs_;

  // Unit: mm
  static constexpr float SMALL_ARMOR_WIDTH = 135;
  static constexpr float SMALL_ARMOR_HEIGHT = 55;
  static constexpr float LARGE_ARMOR_WIDTH = 225;
  static constexpr float LARGE_ARMOR_HEIGHT = 55;

  // Four vertices of armor in 3d
  std::vector<cv::Point3f> small_armor_points_;
  std::vector<cv::Point3f> large_armor_points_;

  cv::Point2f reproject(Eigen::Vector3d &xyz); 
  /**
   * @brief 计算yaw的代价函数
   * 
   * @param yaw 输入的yaw
   * @param tvec 三维平移向量
   * @param object_points 装甲板在装甲板坐标系下的坐标
   * @param image_armor_points 装甲板在图像中的坐标
   * @return double 
   */
  double countRemapYawCost(double yaw, cv::Mat tvec,std::vector<cv::Point3f> object_points,std::vector<cv::Point2f> image_armor_points);
  /**
   * @brief 对yaw重投影，三分搜索以获取更为精确的yaw
   * 
   * @param left 左边界
   * @param right 右边界
   * @param epsilon 精度
   * @param image_armor_points  装甲板在图像中的坐标
   * @param tvec PNP解算的平移向量
   * @param object_points 装甲板在装甲板坐标系的坐标
   * @return double 
   */
  double goldenSreach(double left, double right, double epsilon,std::vector<cv::Point2f> &image_armor_points, cv::Mat tvec,std::vector<cv::Point3f> object_points);
};

}  // namespace rm_auto_aim

#endif  // ARMOR_DETECTOR__PNP_SOLVER_HPP_
